import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.path as path
import numpy as np

# 创建画布
fig, ax = plt.subplots(figsize=(16, 12), dpi=100)

# 设置背景色
ax.set_facecolor('#f8f9fa')

# 定义颜色
colors = {
    'main': '#f9d77e',
    'client': '#a8d1df',
    'processor': '#c3e6cb',
    'store': '#ffeeba',
    'model': '#bee5eb',
    'doc': '#d6d8db',
    'function': '#f5c6cb',
    'flow': '#fdfdfe',
    'arrow': '#555555'
}

# 绘制矩形函数
def draw_box(x, y, width, height, label, color, fontsize=10):
    rect = patches.Rectangle((x, y), width, height, linewidth=1, edgecolor='black', facecolor=color, alpha=0.8)
    ax.add_patch(rect)
    ax.text(x + width/2, y + height/2, label, ha='center', va='center', fontsize=fontsize, fontweight='bold')

# 绘制箭头函数
def draw_arrow(x1, y1, x2, y2, label='', curved=False):
    if curved:
        # 创建弯曲的箭头
        mid_x = (x1 + x2) / 2
        mid_y = min(y1, y2) - 0.5
        verts = [(x1, y1), (mid_x, mid_y), (x2, y2)]
        codes = [path.Path.MOVETO, path.Path.CURVE3, path.Path.CURVE3]
        arrow_path = path.Path(verts, codes)
        arrow_patch = patches.PathPatch(arrow_path, facecolor='none', edgecolor=colors['arrow'], lw=1)
        ax.add_patch(arrow_patch)
        
        # 添加箭头
        dx = x2 - mid_x
        dy = y2 - mid_y
        arrow_length = np.sqrt(dx**2 + dy**2)
        dx, dy = dx/arrow_length, dy/arrow_length
        ax.arrow(mid_x, mid_y, dx*0.9, dy*0.9, head_width=0.1, head_length=0.1, 
                fc=colors['arrow'], ec=colors['arrow'], length_includes_head=True)
    else:
        # 创建直线箭头
        ax.arrow(x1, y1, x2-x1, y2-y1, head_width=0.1, head_length=0.1, 
                fc=colors['arrow'], ec=colors['arrow'], length_includes_head=True)
    
    # 添加标签
    if label:
        ax.text((x1+x2)/2, (y1+y2)/2, label, ha='center', va='center', 
                fontsize=8, bbox=dict(facecolor='white', alpha=0.7, boxstyle='round,pad=0.2'))


# 绘制主要组件
# 1. Main Classes
draw_box(5, 9, 2, 1, 'main()', colors['main'], 12)
draw_box(5, 7, 2, 1, 'RAGSystem', colors['main'], 12)
# 2. Core Components
draw_box(2, 5, 2, 1, 'DeepSeekClient', colors['client'])
draw_box(5, 5, 2, 1, 'DocumentProcessor', colors['processor'])
draw_box(8, 5, 2, 1, 'VectorStore', colors['store'])
# 3. Vector Storage Dependencies
draw_box(8, 3, 2, 1, 'EmbeddingModel', colors['model'])
draw_box(8, 1, 2, 1, 'FAISS Index', colors['store'])
# 4. Embedding Model Options
draw_box(5, 1, 2, 0.6, 'SentenceTransformer', colors['model'])
draw_box(7, 1, 2, 0.6, 'Transformers', colors['model'])
draw_box(9, 1, 2, 0.6, 'TF-IDF', colors['model'])
draw_box(11, 1, 2, 0.6, 'Random Embedding', colors['model'])
# 5. Document Class
draw_box(2, 3, 2, 1, 'Document', colors['doc'])
# 6. RAG System Functions
draw_box(1, 7, 1.5, 0.6, 'Add Document', colors['function'])
draw_box(3, 7, 1.5, 0.6, 'Add Text', colors['function'])
draw_box(10, 7, 1.5, 0.6, 'Query', colors['function'])
draw_box(12, 7, 1.5, 0.6, 'Save/Load', colors['function'])
# 7. Query Flow
draw_box(10, 5.5, 1.5, 0.6, 'Retrieve Documents', colors['flow'])
draw_box(10, 4.5, 1.5, 0.6, 'Build Context', colors['flow'])
draw_box(10, 3.5, 1.5, 0.6, 'Construct Prompt', colors['flow'])
draw_box(10, 2.5, 1.5, 0.6, 'Call API', colors['flow'])
draw_box(10, 1.5, 1.5, 0.6, 'Return Results', colors['flow'])
# Draw Connection Arrows
# Main Connections
draw_arrow(6, 9, 6, 8)
draw_arrow(5, 7, 3, 5.5)
draw_arrow(6, 7, 6, 6)
draw_arrow(7, 7, 9, 5.5)
# Component Connections
draw_arrow(9, 5, 9, 4)
draw_arrow(9, 3, 9, 2)
draw_arrow(8, 3.5, 3, 3.5)
# Embedding Model Selection
draw_arrow(9, 3, 6, 1.6, 'Primary')
draw_arrow(9, 3, 8, 1.6, 'Alternative 1')
draw_arrow(9, 3, 10, 1.6, 'Alternative 2')
draw_arrow(9, 3, 12, 1.6, 'Alternative 3')
# Function Connections
draw_arrow(1.75, 7, 3, 5.5, 'Process Files')
draw_arrow(3.75, 7, 6, 5.5, 'Process Text')
draw_arrow(10.75, 7, 10.75, 6.1)
draw_arrow(10.75, 5.5, 10.75, 5.1)
draw_arrow(10.75, 4.5, 10.75, 4.1)
draw_arrow(10.75, 3.5, 10.75, 3.1)
draw_arrow(10.75, 2.5, 10.75, 2.1)

# 设置图表属性
ax.set_xlim(0, 14)
ax.set_ylim(0, 10)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('DeepSeek RAG System Flow Chart', fontsize=16, fontweight='bold')
# 添加图例
legend_elements = [
    patches.Patch(facecolor=colors['main'], edgecolor='black', label='Main Class'),
    patches.Patch(facecolor=colors['client'], edgecolor='black', label='API Client'),
    patches.Patch(facecolor=colors['processor'], edgecolor='black', label='Document Processor'),
    patches.Patch(facecolor=colors['store'], edgecolor='black', label='Vector Store'),
    patches.Patch(facecolor=colors['model'], edgecolor='black', label='Embedding Model'),
    patches.Patch(facecolor=colors['doc'], edgecolor='black', label='Data Structure'),
    patches.Patch(facecolor=colors['function'], edgecolor='black', label='System Function'),
    patches.Patch(facecolor=colors['flow'], edgecolor='black', label='Query Flow')
]
ax.legend(handles=legend_elements, loc='upper center', bbox_to_anchor=(0.5, -0.05),
          fancybox=True, shadow=True, ncol=4)

# 保存图表
plt.tight_layout()
plt.savefig('work7/deepseek_rag_flowchart.png', dpi=300, bbox_inches='tight')
plt.close()

print("流程图已保存为 'work7/deepseek_rag_flowchart.png'")